import os
import json
import httpx
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

api_key = os.getenv("BAILIAN_KEY")
base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
model = "qwen2.5-7b-instruct-1m"  # "qwen-plus"

http_client = httpx.Client(
    # proxy="http://localhost:8080",
    verify=False,
)

client = OpenAI(
    api_key=api_key,
    base_url=base_url,
    http_client=http_client,
)


# 定义本地函数：加法
def add_numbers(a: int, b: int) -> int:
    """计算两个整数的和"""
    return a + b


# 函数注册表
available_functions = {"add_numbers": add_numbers}

# 系统提示：识别意图并根据函数签名生成参数
system_prompt_intent = r"""
你是一个智能助手，负责处理用户请求，识别意图，提取参数，并指定输出格式。你的任务是：
1. 分析用户输入，确定请求的函数、参数和输出格式。
2. 支持的函数：
   - add_numbers: 计算两个整数的和。参数：a (int), b (int)
3. 支持的输出格式：latex, markdown
4. 根据用户输入和函数签名，生成一个 JSON 对象，结构如下：
{
  "function": "<function_name>",
  "arguments": {
    "<param1>": <value1>,
    "<param2>": <value2>,
    ...
  },
  "output_format": "<format>"
}
5. 确保生成的 parameters 严格符合函数签名：
   - add_numbers 需要两个整数参数：a 和 b。
   - 如果用户输入中缺少参数或类型不正确，返回：
{
  "function": "none",
  "error": "参数缺失或类型错误"
}
6. 仅返回 JSON 对象，不包含额外文本。
"""

# 用户输入
user_input = "计算150和250之和，并将结果按latex输出"

# 第一次调用：识别意图并生成参数
response_intent = client.chat.completions.create(
    model=model,
    messages=[
        {"role": "system", "content": system_prompt_intent},
        {"role": "user", "content": user_input},
    ],
    max_tokens=100,
    temperature=0.0,
)

print(response_intent.choices[0].message)

# 解析第一次响应的 JSON
response_json_intent = json.loads(response_intent.choices[0].message.content.strip())
function_name = response_json_intent.get("function")
arguments = response_json_intent.get("arguments", {})
output_format = response_json_intent.get("output_format", "latex")

# 检查是否识别到有效函数和参数
if function_name == "none":
    print(f"错误：{response_json_intent.get('error', '无法识别请求')}")
    exit()

# 调用本地函数
if function_name in available_functions:
    result = available_functions[function_name](**arguments)
else:
    print("不支持的函数")
    exit()

# 系统提示：格式化输出
system_prompt_format = r"""
你是一个助手，负责根据用户请求格式化函数调用的结果。你的任务是：
1. 接收用户的原始输入、函数结果和指定的输出格式。
2. 根据指定的输出格式生成响应。
3. 确保响应清晰、简洁且格式正确。
- 如果格式是 "latex"，将结果包裹在 LaTeX 数学符号中（例如 \( result \)）。
- 如果格式是 "markdown"，使用适当的 Markdown 语法。
仅返回格式化后的响应，不包含额外说明。
"""

# 第二次调用：生成最终响应
response_format = client.chat.completions.create(
    model=model,
    messages=[
        {"role": "system", "content": system_prompt_format},
        {
            "role": "user",
            "content": f"Original input: {user_input}\nFunction result: {result}\nOutput format: {output_format}",
        },
    ],
    max_tokens=100,
    temperature=0.0,
)

# 输出最终响应
final_response = response_format.choices[0].message.content.strip()
print(final_response)
